from langchain_core.messages import AIMessage
from langgraph.graph import StateGraph, END
from langgraph.types import interrupt, Command
from langgraph.graph.state import CompiledStateGraph
from langchain_core.runnables import RunnableConfig
from states import GraphState
from create_sql_query_code import create_sql_query_code
from execute_sql_database_code import execute_sql_database_code
from fix_sql_database_code import fix_sql_database_code
from langgraph.checkpoint.memory import MemorySaver
import asyncio
import json

AGENT_NAME = "sql_database_agent"

workflow = StateGraph(GraphState)

# Always add create, execute, and fix nodes
workflow.add_node("create_sql_query_code", create_sql_query_code)
workflow.add_node("execute_sql_database_code", execute_sql_database_code)
workflow.add_node("fix_sql_database_code", fix_sql_database_code)

# Conditionally add the recommended-steps node
# if not bypass_recommended_steps:
#     workflow.add_node(recommended_steps_node_name, node_functions[recommended_steps_node_name])

# Conditionally add the human review node
# if human_in_the_loop:
#     workflow.add_node(human_review_node_name, node_functions[human_review_node_name])

# Conditionally add the explanation node
# if not bypass_explain_code:
#     workflow.add_node(explain_code_node_name, node_functions[explain_code_node_name])

# * EDGES

# Set the entry point
# entry_point = create_code_node_name if bypass_recommended_steps else recommended_steps_node_name

workflow.set_entry_point("create_sql_query_code")

# if not bypass_recommended_steps:
#     workflow.add_edge(recommended_steps_node_name, create_code_node_name)

workflow.add_edge("create_sql_query_code", "execute_sql_database_code")
workflow.add_edge("fix_sql_database_code", "execute_sql_database_code")

error_key="sql_database_error"
retry_count_key = "retry_count"
max_retries_key = "max_retries"

# Define a helper to check if we have an error & can still retry
def error_and_can_retry(state):
    return (
        state.get(error_key) is not None
        and state.get(retry_count_key) is not None
        and state.get(max_retries_key) is not None
        and state[retry_count_key] < state[max_retries_key]
    )
    
# If human in the loop, add a branch for human review
# if human_in_the_loop:
#     workflow.add_conditional_edges(
#         execute_code_node_name,
#         lambda s: "fix_code" if error_and_can_retry(s) else "human_review",
#         {
#             "human_review": human_review_node_name,
#             "fix_code": fix_code_node_name,
#         },
#     )
# else:
#     # If no human review, the next node is fix_code if error, else explain_code.
#     if not bypass_explain_code:
#         workflow.add_conditional_edges(
#             execute_code_node_name,
#             lambda s: "fix_code" if error_and_can_retry(s) else "explain_code",
#             {
#                 "fix_code": fix_code_node_name,
#                 "explain_code": explain_code_node_name,
#             },
#         )
#     else:
#         workflow.add_conditional_edges(
#             execute_code_node_name,
#             lambda s: "fix_code" if error_and_can_retry(s) else "END",
#             {
#                 "fix_code": fix_code_node_name,
#                 "END": END,
#             },
#         )

workflow.add_conditional_edges(
    "execute_sql_database_code",
    lambda s: "fix_code" if error_and_can_retry(s) else "END",
    {
        "fix_code": "fix_sql_database_code",
        "END": END,
    },
)

# if not bypass_explain_code:
#     workflow.add_edge(explain_code_node_name, END)

checkpointer = MemorySaver()

# Finally, compile
app = workflow.compile(
    checkpointer=checkpointer,
    name=AGENT_NAME,
)

config = {"configurable": {"thread_id": "user123"}}
async def handle_question():
    input = "列出所有表名"
    input = "查询系统有多少用户"
    input = "用户wanghh2003的邮箱是多少"
    input = "查询用户username=wanghh2003的帐户余额是多少？"
    response = await app.ainvoke(
        {
            "user_instructions": input,
            "max_retries": 3,
            "retry_count": 2
        },
        config=config
    )
    print("---------------------------------")
    print(json.dumps(response, indent=2, ensure_ascii=False))
    # print(response)

asyncio.run(handle_question())

